package com.ruoyi.common.utils.shell.jsch;

import com.jcraft.jsch.*;
import lombok.extern.slf4j.Slf4j;

import java.io.BufferedReader;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Map;
import java.util.Properties;

/**
 * ssh操作
 *
 * @author Hack2003
 * @date 2024/3/26 16:15
 **/
@Slf4j
public class SshUtil {
    /**
     * 建立存放session的map，不同主机端口和用户，它的session唯一
     */
    private static Map<String, Session> sessions = new java.util.concurrent.ConcurrentHashMap<>();

    /**
     * 构建Session的唯一标识键
     *
     * @param host
     * @param port
     * @param user
     * @return
     */
    private static String buildSessionKey(String host, int port, String user) {
        return host + ":" + port + ":" + user;
    }

    /**
     * 获取ssh配置
     *
     * @return
     */
    public static Properties getSshConfig() {
        Properties sshConfig = new Properties();
        //当 StrictHostKeyChecking 设置为 yes （默认值）时，SSH客户端会在~/.ssh/known_hosts文件中查找匹配的主机密钥。如果找不到或者发现服务器提供的密钥与已知密钥不同，SSH客户端将拒绝连接并提示用户确认新的主机密钥。这有助于防止中间人攻击（MITM），因为用户会被警告任何不一致的情况。
        //当 StrictHostKeyChecking 设置为 no 时，SSH客户端不会进行严格的检查，即使服务器的密钥发生变化也不会阻止连接，而是接受任何提供的密钥。这种设置虽然方便，但大大降低了安全性，因为它允许任意服务器冒充已知的主机而无需用户进一步确认。
        sshConfig.setProperty("StrictHostKeyChecking", "no");
        return sshConfig;
    }

    /**
     * 方法负责建立或获取已经存在的Session，使用组合主机名、端口和用户名作为唯一标识
     *
     * @param host
     * @param port
     * @param user
     * @param password
     * @param config
     * @return
     * @throws JSchException
     */
    public static synchronized Session getSession(String host, int port, String user, String password, Properties config) throws JSchException {
        String key = buildSessionKey(host, port, user);
        Session session = sessions.get(key);
        if (session == null || !session.isConnected()) {
            JSch jsch = JschSingleton.getJSchInstance();
            session = jsch.getSession(user, host, port);
            session.setPassword(password);
            session.setConfig(config);
            session.connect();
            sessions.put(key, session);
        }
        return session;
    }
    public static synchronized Session getSession(String host, int port, String user, String password) throws JSchException {
        String key = buildSessionKey(host, port, user);
        Session session = sessions.get(key);
        if (session == null || !session.isConnected()) {
            JSch jsch = JschSingleton.getJSchInstance();
            session = jsch.getSession(user, host, port);
            session.setPassword(password);
            session.setConfig(getSshConfig());
            session.connect();
            sessions.put(key, session);
        }
        return session;
    }
    public static void main(String[] args) {
        String host = "your_host";
        int port = 22;
        String user = "your_user";
        String password = "your_password";
        try {
            //检验模式
            Properties config = getSshConfig();
            //建立session
            Session session = getSession(host, port, user, password, config);
            String command = "your_command";
            //执行命令
            SshExecutionResult result = executeCommand(session, command);
            // 处理结果...
        } catch (JSchException e) {
            // 处理连接失败异常
        } finally {
            closeSession(host, port, user);
        }
    }

    /**
     * 执行远程命令
     *
     * @param session
     * @param command
     * @return
     */
    public static SshExecutionResult executeCommand(Session session, String command) throws JSchException{
        //建立的命令执行通道
        Channel channel = null;
        try {
            // 打开一个通道用于执行命令
            channel = session.openChannel("exec");
            ((ChannelExec) channel).setCommand(command);
            // 启动命令执行
            channel.connect();

            StringBuilder output = new StringBuilder();
            StringBuilder errorOutput = new StringBuilder();
            try (
                    // 获取输入流和错误流
                    InputStream in = channel.getInputStream();
                    InputStream err = ((ChannelExec) channel).getErrStream();
                    // 读取命令执行结果
                    BufferedReader reader = new BufferedReader(new InputStreamReader(in));
                    BufferedReader errorReader = new BufferedReader(new InputStreamReader(err))
            ) {
                String line;
                while ((line = reader.readLine()) != null) {
                    output.append(line).append("\n");
                }
                while ((line = errorReader.readLine()) != null) {
                    if(!line.isEmpty()&&!line.equals("mysqldump: [Warning] Using a password on the command line interface can be insecure.")) {
                        errorOutput.append(line).append("\n");
                        // 如果有错误输出，可根据需要抛出异常或记录日志
                        if (!errorOutput.toString().isEmpty() && line.indexOf("error")!=-1) {
                            throw new JSchException("Error occurred while executing remote command: " + errorOutput.toString());
                        }
                    }
                }
            }
            return new SshExecutionResult(output.toString(), errorOutput.toString());
        } catch (Exception e) {
            log.error("Error occurred while executing remote command: {}", e);
            throw new JSchException("执行远程命令时发生IO错误", e);
        } finally {
            // Channel关闭
            if (channel != null && channel.isConnected()) {
                try {
                    channel.disconnect();
                } catch (Exception ex) {
                    log.error("Failed to close channel.", ex);
                }
            }
        }
    }

    /**
     * ssh执行结果
     */
    public static class SshExecutionResult {
        private final String output;
        private final String errorOutput;

        public SshExecutionResult(String output, String errorOutput) {
            this.output = output;
            this.errorOutput = errorOutput;
        }

        public String getOutput() {
            return output;
        }

        public String getErrorOutput() {
            return errorOutput;
        }
    }

    /**
     * 当不再需要某个连接时，关闭对应的Session
     *
     * @param host
     * @param port
     * @param user
     */
    public static synchronized void closeSession(String host, int port, String user) {
        String key = buildSessionKey(host, port, user);
        Session session = sessions.remove(key);
        if (session != null && session.isConnected()) {
            session.disconnect();
        }
    }
}
